package be.neo4reactome.importer;

import be.neo4reactome.parser.Interaction;
import be.neo4reactome.parser.InteractionsParser;
import be.neo4reactome.parser.Pathway;
import be.neo4reactome.parser.PathwaysParser;
import org.neo4j.graphdb.*;
import org.neo4j.graphdb.index.*;

import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * This class handles the import of the Reactome data into the graph database.
 */
public class Importer {

    private GraphDatabaseService graphDb;
    private Node firstNode;
    private Node secondNode;

    // Index for the protein nodes.
    private Index<Node> proteinIndex;

    // Index for the pathway nodes.
    private Index<Node> pathwayIndex;

    // Index for the interaction nodes.
    private RelationshipIndex interactionIndex;

    // Key names
    public static final String PROTEIN_ACCESSION = "accession";
    public static final String INTERACTION = "interaction";
    public static final String PATHWAY_DESCRIPTION = "pathdesc";

    private String interactionsFile;
    private String pathwayFile;

    // Interaction type
    public static enum InteractionType implements RelationshipType {
        DIRECT_COMPLEX, INDIRECT_COMPLEX, REACTION, NEIGHBOURING_REACTION
    }

    // Protein pathway type
    public static enum ProteinPathwayType implements RelationshipType {
        INVOLVED_IN
    }

    /**
     * Constructer for importing data into the graph database.
     * @param graphDb The graph database system.
     */
    public Importer(GraphDatabaseService graphDb, String interactionsFile, String pathwayFile) {
        this.graphDb = graphDb;
        this.interactionsFile = interactionsFile;
        this.pathwayFile = pathwayFile;

        // Setup the indices
        setupIndices();

        // Create interactions graph
        System.out.println("Starting to setup the interactions graph");
        createInteractionsGraph();

        // Create pathway graph
        System.out.println("Starting to setup the pathways graph");
        createPathwaysGraph();
    }


    /**
     * Returns the interaction type for a specified interaction string.
     * @param interaction
     * @return
     */
    private static InteractionType getInteractionType(String interaction) {
        if(interaction.equalsIgnoreCase("reaction")) {
            return InteractionType.REACTION;
        } else if(interaction.equalsIgnoreCase("neighbouring_reaction")) {
            return InteractionType.NEIGHBOURING_REACTION;
        } else if(interaction.equalsIgnoreCase("direct_complex")) {
            return InteractionType.DIRECT_COMPLEX;
        } else if(interaction.equalsIgnoreCase("indirect_complex")) {
            return InteractionType.INDIRECT_COMPLEX;
        } else {
            // Default interaction type == reaction
            return InteractionType.REACTION;
        }
    }

    /**
     * This method creates the interactions graph.
     */
    private void createInteractionsGraph() {
        List<Interaction> interactions = InteractionsParser.read(interactionsFile);
        Transaction tx = graphDb.beginTx();
        try {
            // Iterate the interactions and add the data to the graphDB.
            for (Interaction interaction : interactions) {
                firstNode = createAndIndexProteinNode(interaction.getProteinA());
                secondNode = createAndIndexProteinNode(interaction.getProteinB());
                createRelationship(firstNode, secondNode, getInteractionType(interaction.getInteraction()));
            }
            // Transaction
            tx.success();
        } finally {
            tx.finish();
        }
    }

    /**
     * This method creates the pathways graph.
     */
    public void createPathwaysGraph() {
        // Get the pathways
        Map<String, List<Pathway>> pathways = PathwaysParser.read(pathwayFile);

        Transaction tx = graphDb.beginTx();
        try {

            // Iterate the entries
            Set<Map.Entry<String, List<Pathway>>> entries = pathways.entrySet();
            for (Map.Entry<String, List<Pathway>> entry : entries) {
                Node protein = createAndIndexProteinNode(entry.getKey());
                List<Pathway> pathwayList = entry.getValue();
                for (Pathway p : pathwayList) {
                    Node pathway = createAndIndexPathwayNode(p.getDescription());
                    createRelationship(protein, pathway, ProteinPathwayType.INVOLVED_IN);
                }
            }
            // Transaction
            tx.success();
        } finally {
            tx.finish();
        }
    }

    /**
     * This method sets up the indices for the graph database.
     */
    private void setupIndices() {
        final IndexManager index = graphDb.index();
        proteinIndex = index.forNodes("proteins");
        interactionIndex = index.forRelationships("interactions");
        pathwayIndex = index.forNodes("pathways");
    }





    /**
     * This method creates and indexes a protein node.
     * @param proteinAccession The protein accession (identifier for a protein).
     * @return (protein) node
     */
    private Node createAndIndexProteinNode(final String proteinAccession) {
        return new UniqueFactory.UniqueNodeFactory(proteinIndex) {
            protected void initialize(Node node, Map<String, Object> props) {
                node.setProperty(PROTEIN_ACCESSION, proteinAccession);
            }
        }.getOrCreate(PROTEIN_ACCESSION, proteinAccession);
    }


    /**
     * This method creates and indexes a pathway node.
     * @param pathwayDescription The pathway description (identifier for a pathway).
     * @return (pathway) node
     */
    private Node createAndIndexPathwayNode(final String pathwayDescription) {
        Node pathwayNode = pathwayIndex.get(PATHWAY_DESCRIPTION, pathwayDescription).getSingle();

        if (pathwayNode != null) return pathwayNode;

        Node node = graphDb.createNode();
        node.setProperty(PATHWAY_DESCRIPTION, pathwayDescription);
        pathwayIndex.add(node, PATHWAY_DESCRIPTION, pathwayDescription);
        return node;
    }

    /**
     * Creates an relationship.
     * @param proteinA The first node.
     * @param proteinB The second node.
     * @param interactionName The relationship type.
     */
    private void createRelationship(Node proteinA, Node proteinB, RelationshipType interactionName) {
        final String name = interactionName.name();

        Relationship interaction = interactionIndex.get(INTERACTION, name,proteinA,proteinB).getSingle();

        if (interaction==null) {
            interaction = proteinA.createRelationshipTo(proteinB, interactionName);
            interaction.setProperty(INTERACTION, name);
            interactionIndex.add(interaction,INTERACTION, name);
        }
    }
}
